import pandas as pd
import numpy as np


def filter_admet_rows(admet_df):
    """
    根据条件 "Caco-2" + "CYP3A4" + 1 - "hERG" + "HOB" + 1 - "MN" >= 3
    筛选符合条件的行号。

    Parameters:
    - admet_df: 包含 ADMET 数据的 DataFrame

    Returns:
    - 符合条件的行号列表
    """
    # 计算符合条件的行
    condition = (
        admet_df["Caco-2"]
        + admet_df["CYP3A4"]
        + (1 - admet_df["hERG"])
        + admet_df["HOB"]
        + (1 - admet_df["MN"])
    ) >= 3

    # 返回符合条件的行号
    return admet_df.index[condition].tolist()


if __name__ == "__main__":
    # 加载ADMET数据
    admet_df = pd.read_csv("data/ADMET_training.csv")
    ER_activity_df = pd.read_csv("data/ER_activity_training.csv")
    # 筛选符合条件的行号
    matching_rows = filter_admet_rows(admet_df)

    filtered_er_activity_df = ER_activity_df.loc[matching_rows]

    # 按 pic50 属性值降序排序并选出前 50 个
    top_50_rows = filtered_er_activity_df.nlargest(50, "pIC50")

    top_50_row_indices = top_50_rows.index.tolist()

    print(top_50_row_indices)
